import json
import random
from collections import defaultdict
import sys



data_p = sys.argv[1]


def get_source(item):
    """Identify the source of an item from possible keys."""
    if "datasource" in item:
        return item["datasource"]
    elif "source" in item:
        return item["source"]
    elif "metadata" in item and isinstance(item["metadata"], dict):
        return item["metadata"].get("dataset")
    return 'default'  # No source identified


# Load the JSON data
with open(data_p, "r") as f:
    data = json.load(f)

# Group items by their source
groups = defaultdict(list)

video_num = 0
for item in data:
    source = get_source(item)
    # print(len(item['conversations']))
    if 'image' in item and isinstance(item['image'], list) and len(item['image']) > 1:
        # print(f'{item} image more than 1..')
        pass
    if len(item['conversations']) > 2:
        continue
        # is_multiple_image = False
        # for idx, conv in enumerate(item['conversations']):
        #     if idx > 1 and '<image>' in conv['value']:
        #         is_multiple_image = True
                
      
        # if is_multiple_image:
        #     print(f'{item} have multiple conversations!')
        #     continue
    if 'image' not in item:
        # print(f'pure text: {item}, {len(item['conversations'][1]['value'])}')
        continue
    if 'video' in item:
        video_num += 1
    groups[source].append(item)

print(f'all: {len(data)} video: {video_num}')
# Seed for reproducibility
random.seed(42)
filtered_data = []

print(groups.keys())

# Process each group
for source, items in groups.items():
    # Shuffle to ensure random selection
    if 'nuscenes' in source or 'RAVEN' in source:
        filtered_data.extend(items)
    else:
        random.shuffle(items)

        # Calculate 50% of the group (at least 1 item if not empty)
        group_size = len(items)
        k = max(1, int(group_size * 1)) if group_size > 0 else 0

        print(f'{source} {len(items)} {(len(items)/len(data))*100:.2f}%')
        # Select items and add to result
        filtered_data.extend(items[:k])
        
# for d in data:
#     if 'image' in d:
#         filtered_data.append(d)
#     else:
#         print(f'{d} not have image ')

# Save the filtered data
sub_d = data_p.replace(".json", f"_sub_{len(filtered_data)}.json")
with open(sub_d, "w") as f:
    json.dump(filtered_data, f, indent=2, ensure_ascii=False)
print(f'saved into: {sub_d}')
